"""
    包含：
    - 基本的 Pool.map 示例
    - Pool.apply_async 示例（带回调）
    - Pool.imap_unordered 示例（带超时与分块）
    - 使用上下文管理器或 close()/join()
    - maxtasksperchild 的使用
    - 捕获 KeyboardInterrupt 正确退出进程池
    - 统一使用上下文管理（with Pool(...)）以确保资源正确释放
    - 在主进程中捕获 KeyboardInterrupt 并优雅终止子进程
    - 增加 error_callback，避免 apply_async 异常丢失
    - 自动计算合理的 chunksize，降低进程间通信开销
    - 可设置 maxtasksperchild 防止子进程内存泄露累积
    - 在 macOS 上尽量使用 spawn 启动方式以避免子进程重复执行主逻辑
"""
from multiprocessing import Pool, cpu_count, current_process
import multiprocessing as mp
import time
import os
import signal
import sys
import traceback

def _init_worker():
    """
    子进程初始化函数：
    - 忽略 SIGINT（主进程负责处理 Ctrl+C），避免子进程被中断信号直接杀死。
    - 可用于初始化子进程级资源（数据库连接、随机种子等）。
    """
    signal.signal(signal.SIGINT, signal.SIG_IGN)

def heavy_task(x):
    """
    模拟耗时计算的工作函数。
    返回包含当前进程 pid 的结果，便于观察进程池分配情况。
    """
    pid = os.getpid()
    proc_name = current_process().name
    time.sleep(0.5)  # 模拟耗时
    return f"input={x} pid={pid} proc={proc_name} result={x*x}"

def _calc_chunksize(n_items, pool_size, target_chunks_per_process=4):
    """
    根据任务数与进程数计算合理的 chunksize：
    - 避免 chunksize 过小导致通信/调度开销大
    - 避免 chunksize 过大导致负载不均衡
    """
    if n_items <= 0 or pool_size <= 0:
        return 1
    return max(1, n_items // (pool_size * target_chunks_per_process))

def safe_map(func, inputs, pool_size=None):
    """安全封装 Pool.map，支持 KeyboardInterrupt 优雅退出。"""
    if pool_size is None:
        pool_size = min(4, cpu_count())
    chunksize = _calc_chunksize(len(inputs), pool_size)
    try:
        with Pool(processes=pool_size, initializer=_init_worker) as pool:
            return pool.map(func, inputs, chunksize=chunksize)
    except KeyboardInterrupt:
        print("主进程收到 KeyboardInterrupt，已中止 map 操作。")
        raise

def safe_apply_async(func, inputs, pool_size=None, on_result=None, on_error=None, maxtasksperchild=10):
    """
    使用 apply_async 提交多个任务并等待完成。
    - on_result(result) 在主进程调用（每个任务完成时）
    - on_error(exc_info_tuple) 在主进程调用（任务异常时）
    """
    if pool_size is None:
        pool_size = min(3, cpu_count())

    def _error_callback(e):
        exc_info = (type(e).__name__, str(e), traceback.format_exc())
        if on_error:
            try:
                on_error(exc_info)
            except Exception:
                print("on_error 回调执行出错：", traceback.format_exc())
        else:
            print("task error:", exc_info)

    pool = None
    try:
        pool = Pool(processes=pool_size, initializer=_init_worker, maxtasksperchild=maxtasksperchild)
        for x in inputs:
            pool.apply_async(func, args=(x,), callback=on_result, error_callback=_error_callback)
        pool.close()
        pool.join()
    except KeyboardInterrupt:
        print("主进程收到 KeyboardInterrupt，正在终止子进程...")
        if pool is not None:
            try:
                pool.terminate()
                pool.join()
            except Exception:
                pass
        raise

def imap_unordered_example():
    """
    imap_unordered 返回结果迭代器，任务完成即可获取结果（不保证输入顺序）。
    使用 chunksize 控制调度粒度。
    """
    inputs = list(range(12))
    pool_size = min(4, cpu_count())
    chunksize = _calc_chunksize(len(inputs), pool_size)
    try:
        with Pool(processes=pool_size, initializer=_init_worker) as pool:
            for res in pool.imap_unordered(heavy_task, inputs, chunksize=chunksize):
                print("imap_unordered got:", res)
    except KeyboardInterrupt:
        print("imap_unordered 被中断，退出。")
        raise

def timeout_example():
    """
    对单个 apply_async 使用 get(timeout=...) 演示超时控制。
    超时发生后通常需要根据业务决定是否 terminate 池。
    """
    pool_size = 2
    with Pool(processes=pool_size, initializer=_init_worker) as pool:
        try:
            async_res = pool.apply_async(heavy_task, args=(9,))
            try:
                # 故意设置短超时以示例超时处理
                print("result with timeout:", async_res.get(timeout=0.1))
            except Exception as e:
                print("get 超时或出错:", type(e).__name__, e)
                print("由于超时，终止池（避免子进程继续运行）...")
                pool.terminate()
        finally:
            try:
                pool.join()
            except Exception:
                pass

def simple_map_example():
    """使用 safe_map 演示并行 map（返回按输入顺序的结果）。"""
    inputs = list(range(8))
    pool_size = min(4, cpu_count())
    results = safe_map(heavy_task, inputs, pool_size=pool_size)
    print("map results:", results)

def apply_async_example():
    """使用 safe_apply_async 示例（带回调与错误回调）。"""
    def on_result(res):
        print("callback got:", res)

    def on_error(exc_info):
        print("callback error:", exc_info)

    inputs = list(range(6))
    safe_apply_async(heavy_task, inputs, pool_size=min(3, cpu_count()), on_result=on_result, on_error=on_error)

def main():
    # 尝试设置启动方式（macOS 上 spawn 更稳健），如果已设置则忽略 RuntimeError
    try:
        mp.set_start_method("spawn", force=False)
    except RuntimeError:
        pass

    print("CPU cores:", cpu_count())
    try:
        print("示例：Pool.map")
        simple_map_example()
        print("\n示例：Pool.apply_async（带回调）")
        apply_async_example()
        print("\n示例：Pool.imap_unordered（流式输出）")
        imap_unordered_example()
        print("\n示例：超时处理")
        timeout_example()
        print("\n全部示例执行结束。")
    except KeyboardInterrupt:
        print("主进程被中断 (KeyboardInterrupt)，已退出。")
        sys.exit(1)

if __name__ == "__main__":
    # 必须在主模块保护下运行 multiprocessing 入口
    main()

